import os
import time

import torch
from torch import argmax
from torch.utils.data import DataLoader
from torchvision.transforms import Compose, ToTensor

from dataset import MnistDataset
from model import MLP


if __name__ == "__main__":
    # pred
    output_dir="output/20240403-113712"
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")


    train_img_dir = "./data/MNIST/train"
    train_label_txt = "./data/MNIST/train.txt"
    test_img_dir = "./data/MNIST/test"
    test_label_txt = "./data/MNIST/test.txt"
    epoch_num = 3
    batchsize = 64
    output_root_dir = "./output"

    transform = Compose([ToTensor()])

    test_dataset = MnistDataset(img_dir_path=test_img_dir, label_txt_path=test_label_txt, transform=transform)
    test_dl = DataLoader(test_dataset, batch_size=batchsize, shuffle=False, num_workers=4)

    model = MLP(1).to(device)
    model_path_now = os.path.join(output_dir, f"model_end.pth")
    model.load_state_dict(torch.load(model_path_now, map_location=device))
    print(f"Model weights loaded from {model_path_now}")
    model.eval()
    with torch.no_grad():
        _, (imgs, labels) = next(enumerate(test_dl))
        imgs, labels = imgs.to(device), labels.to(device)
        # print(imgs.shape, labels.shape)
        output = model(imgs)
        # print(output.shape)
        pred = argmax(output, dim=1)
        # print(pred.shape)

        print(f"pred:{pred}\n label:{labels}")
